#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
LangGraph 完整示例 - 使用真实的API
这个示例展示了如何使用LangGraph创建一个完整的Agent
"""

import os
from typing import List, Dict, Any
from langchain_core.tools import tool
from langchain_core.messages import HumanMessage, AIMessage
import json

# 尝试加载环境变量（如果有dotenv的话）
try:
    from dotenv import load_dotenv
    load_dotenv()
except ImportError:
    print("💡 提示：安装python-dotenv包可以更方便地管理环境变量")
    print("   pip install python-dotenv")

# 定义工具函数
@tool
def get_weather(city: str) -> str:
    """获取指定城市的天气信息"""
    # 在实际应用中，这里会调用真实的天气API
    weather_data = {
        "北京": "晴天，温度25°C，湿度60%，风速3m/s",
        "上海": "多云，温度22°C，湿度70%，风速2m/s", 
        "广州": "雨天，温度28°C，湿度85%，风速1m/s",
        "深圳": "晴天，温度30°C，湿度65%，风速2m/s",
        "杭州": "阴天，温度20°C，湿度75%，风速1m/s",
        "成都": "多云，温度24°C，湿度68%，风速2m/s",
        "西安": "晴天，温度26°C，湿度55%，风速3m/s",
        "武汉": "多云，温度23°C，湿度72%，风速2m/s",
        "重庆": "雨天，温度27°C，湿度80%，风速1m/s",
        "天津": "晴天，温度24°C，湿度58%，风速2m/s"
    }
    return weather_data.get(city, f"{city}的天气：晴天，温度适宜，湿度正常")

@tool
def get_time_info() -> str:
    """获取当前时间信息"""
    import datetime
    now = datetime.datetime.now()
    weekday = ["星期一", "星期二", "星期三", "星期四", "星期五", "星期六", "星期日"][now.weekday()]
    return f"当前时间：{now.strftime('%Y年%m月%d日')} {weekday} {now.strftime('%H:%M:%S')}"

@tool
def calculate_math(expression: str) -> str:
    """计算数学表达式"""
    try:
        # 安全的数学计算
        allowed_chars = set('0123456789+-*/.() ')
        if not all(c in allowed_chars for c in expression):
            return "错误：只支持基本的数学运算符（+, -, *, /, ()）"
        
        result = eval(expression)
        return f"计算结果：{expression} = {result}"
    except Exception as e:
        return f"计算错误：{str(e)}"

@tool
def get_city_info(city: str) -> str:
    """获取城市基本信息"""
    city_info = {
        "北京": "中国首都，人口约2200万，著名景点：故宫、长城、天坛",
        "上海": "中国经济中心，人口约2400万，著名景点：外滩、东方明珠、豫园",
        "广州": "广东省省会，人口约1500万，著名景点：广州塔、陈家祠、白云山",
        "深圳": "经济特区，人口约1300万，著名景点：世界之窗、深圳湾公园、大梅沙",
        "杭州": "浙江省省会，人口约1200万，著名景点：西湖、灵隐寺、钱塘江",
        "成都": "四川省省会，人口约2100万，著名景点：宽窄巷子、大熊猫基地、武侯祠"
    }
    return city_info.get(city, f"抱歉，我还没有{city}的详细信息")

# 检查是否有API密钥
def check_api_availability():
    """检查API可用性"""
    openai_key = os.getenv("OPENAI_API_KEY")
    anthropic_key = os.getenv("ANTHROPIC_API_KEY")
    
    if openai_key:
        return "openai", openai_key
    elif anthropic_key:
        return "anthropic", anthropic_key
    else:
        return "none", None

def create_agent_with_api():
    """创建带API的Agent"""
    provider, api_key = check_api_availability()
    
    if provider == "openai":
        from langchain_openai import ChatOpenAI
        from langgraph.prebuilt import create_react_agent
        from langchain_core.messages import SystemMessage
        
        # 创建模型实例
        model = ChatOpenAI(
            model="gpt-3.5-turbo",
            temperature=0.7
        )
        
        # 创建系统消息作为状态修改器
        system_message = SystemMessage(content="""您是一个友好的智能助手。您可以：
1. 查询天气信息
2. 获取时间信息
3. 进行数学计算
4. 提供城市信息

请根据用户的问题选择合适的工具，并用中文回答。""")
        
        # 创建Agent，使用state_modifier替代system_prompt
        agent = create_react_agent(
            model=model,
            tools=[get_weather, get_time_info, calculate_math, get_city_info],
            state_modifier=system_message
        )
        
        return agent, provider
    
    elif provider == "anthropic":
        from langchain_anthropic import ChatAnthropic
        from langgraph.prebuilt import create_react_agent
        from langchain_core.messages import SystemMessage
        
        # 创建模型实例
        model = ChatAnthropic(
            model="claude-3-sonnet-20240229",
            temperature=0.7
        )
        
        # 创建系统消息
        system_message = SystemMessage(content="""您是一个友好的智能助手。您可以：
1. 查询天气信息
2. 获取时间信息
3. 进行数学计算
4. 提供城市信息

请根据用户的问题选择合适的工具，并用中文回答。""")
        
        # 创建Agent
        agent = create_react_agent(
            model=model,
            tools=[get_weather, get_time_info, calculate_math, get_city_info],
            state_modifier=system_message
        )
        
        return agent, provider
    
    else:
        return None, "none"

# 简单的本地Agent（无需API）
class LocalAgent:
    def __init__(self):
        self.tools = {
            "get_weather": get_weather,
            "get_time_info": get_time_info,
            "calculate_math": calculate_math,
            "get_city_info": get_city_info
        }
        
    def _analyze_input(self, user_input: str) -> Dict[str, Any]:
        """分析用户输入"""
        user_input_lower = user_input.lower()
        
        # 天气询问
        if any(word in user_input_lower for word in ["天气", "weather", "温度", "下雨", "晴天", "阴天"]):
            cities = ["北京", "上海", "广州", "深圳", "杭州", "成都", "西安", "武汉", "重庆", "天津"]
            for city in cities:
                if city in user_input:
                    return {"tool": "get_weather", "params": {"city": city}}
            return {"tool": "get_weather", "params": {"city": "北京"}}
        
        # 时间询问
        elif any(word in user_input_lower for word in ["时间", "几点", "现在", "time", "日期"]):
            return {"tool": "get_time_info", "params": {}}
        
        # 数学计算
        elif any(word in user_input_lower for word in ["计算", "算", "等于"]) or any(char in user_input for char in "+-*/="):
            import re
            math_pattern = r'[\d+\-*/\.\(\)\s]+'
            match = re.search(math_pattern, user_input)
            if match:
                expression = match.group().strip()
                return {"tool": "calculate_math", "params": {"expression": expression}}
        
        # 城市信息
        elif any(word in user_input_lower for word in ["城市", "地方", "景点", "信息", "介绍"]):
            cities = ["北京", "上海", "广州", "深圳", "杭州", "成都"]
            for city in cities:
                if city in user_input:
                    return {"tool": "get_city_info", "params": {"city": city}}
        
        return {"tool": "chat", "params": {}}
    
    def chat(self, user_input: str) -> str:
        """处理用户输入"""
        analysis = self._analyze_input(user_input)
        
        if analysis["tool"] == "chat":
            return ("我是一个智能助手，可以帮您：\n"
                    "1. 查询天气：'北京天气怎么样？'\n"
                    "2. 查询时间：'现在几点？'\n"
                    "3. 数学计算：'计算 2+3*4'\n"
                    "4. 城市信息：'介绍一下上海'\n"
                    "请告诉我您需要什么帮助！")
        
        tool_name = analysis["tool"]
        params = analysis["params"]
        
        if tool_name in self.tools:
            try:
                result = self.tools[tool_name].invoke(params)
                return f"为您查询到：\n{result}"
            except Exception as e:
                return f"查询出错：{str(e)}"
        
        return "抱歉，我无法处理这个请求。"

def main():
    """主函数"""
    print("=" * 60)
    print("🤖 LangGraph 完整示例")
    print("=" * 60)
    
    # 检查API可用性
    provider, api_key = check_api_availability()
    
    if provider != "none":
        print(f"✅ 检测到 {provider.upper()} API密钥，将使用完整的LangGraph Agent")
        try:
            agent, provider_name = create_agent_with_api()
            print(f"🚀 使用 {provider_name.upper()} 模型创建Agent成功")
            agent_type = "api"
        except Exception as e:
            print(f"❌ API Agent创建失败：{str(e)}")
            print("🔄 切换到本地模拟Agent")
            agent = LocalAgent()
            agent_type = "local"
    else:
        print("ℹ️  未检测到API密钥，使用本地模拟Agent")
        print("💡 要使用完整功能，请在.env文件中设置OPENAI_API_KEY或ANTHROPIC_API_KEY")
        agent = LocalAgent()
        agent_type = "local"
    
    print("\n支持的功能：")
    print("1. 天气查询：'北京天气怎么样？'")
    print("2. 时间查询：'现在几点？'")
    print("3. 数学计算：'计算 15 + 25 * 2'")
    print("4. 城市信息：'介绍一下上海'")
    print("5. 输入 'quit' 或 'exit' 退出")
    print("=" * 60)
    
    # 演示示例
    print("\n📝 演示示例：")
    demo_questions = [
        "上海天气怎么样？",
        "现在几点了？",
        "计算 12 * 8 + 5",
        "介绍一下杭州"
    ]
    
    for question in demo_questions:
        print(f"\n示例问题：{question}")
        
        if agent_type == "api":
            try:
                response = agent.invoke({
                    "messages": [HumanMessage(content=question)]
                })
                
                # 提取最后一条消息
                last_message = response["messages"][-1]
                answer = last_message.content if hasattr(last_message, 'content') else str(last_message)
                print(f"回答：{answer}")
            except Exception as e:
                print(f"API调用错误：{str(e)}")
        else:
            answer = agent.chat(question)
            print(f"回答：{answer}")
        
        print("-" * 40)
    
    # 交互式对话
    print("\n🎯 现在您可以开始提问了！")
    
    while True:
        try:
            user_input = input("\n您的问题：").strip()
            
            if user_input.lower() in ['quit', 'exit', '退出', 'bye', 'q']:
                print("👋 再见！感谢使用LangGraph助手！")
                break
            
            if not user_input:
                print("请输入有效的问题。")
                continue
            
            if agent_type == "api":
                try:
                    response = agent.invoke({
                        "messages": [HumanMessage(content=user_input)]
                    })
                    
                    # 提取回答
                    last_message = response["messages"][-1]
                    answer = last_message.content if hasattr(last_message, 'content') else str(last_message)
                    print(f"🤖 助手：{answer}")
                except Exception as e:
                    print(f"❌ 处理出错：{str(e)}")
            else:
                answer = agent.chat(user_input)
                print(f"🤖 助手：{answer}")
                
        except KeyboardInterrupt:
            print("\n\n👋 程序已终止。再见！")
            break
        except Exception as e:
            print(f"❌ 发生错误：{str(e)}")
            print("请重试...")

if __name__ == "__main__":
    main() 